#include "pch.h"
#include "KMeans.h"
#include <vector>
#include <vector>


KMeans::KMeans(uint cluster, uint max_iterator, float min_dis): _cluster(cluster), _max_iterator(max_iterator),
                                                                _min_dis(min_dis)
{
	
}

void KMeans::fit(const std::vector<Vec>& points)
{
	std::vector<uint> indexs;
	float last_dis = FLT_MAX;
	_center = initCenter(points,_cluster);
	for(uint i = 0 ; i< _max_iterator;i ++)
	{
		float dis;
		calDis(_center, points, indexs, dis,_cluster);
		updateCenter(points, indexs);
		if(dis < _min_dis)
			break;
	}
}

std::vector<Cluster::uint> KMeans::predict(const std::vector<Vec>& points) const
{
	std::vector<uint> res; float dis;
	calDis(_center, points, res, dis,_cluster);
	return res;
}

void KMeans::updateCenter(const std::vector<Vec>& points, const std::vector<uint>& indexs)
{
	assert(points.size() == indexs.size());

	for (auto& p : _center)
	{
		p.x = 0; p.y = 0;
	}
	std::vector<uint> sizes(_cluster, 0);

	for (uint i = 0; i < points.size(); i++)
	{
		_center[indexs[i]] += points[i];
		++sizes[indexs[i]];
	}

	for (uint i = 0; i < _cluster; i++)
	{
		_center[i].x /= sizes[i];
		_center[i].y /= sizes[i];
	}
}